package backend;

import backend.value.*;
import middleware.llvm.BasicBlock;
import middleware.llvm.enums.IRType;
import middleware.llvm.instruction.*;
import middleware.llvm.value.ArrayIRType;
import middleware.llvm.value.FunctionIRType;
import middleware.llvm.value.ImmediateValue;
import middleware.llvm.value.Value;
import middleware.optimize.ConflictDiagramBuilder;
import middleware.optimize.LiveVariableAnalyzer;

import java.util.*;
import java.util.stream.IntStream;

/**
 * 值管理器，用于管理目标代码中的值
 *
 * @author 杨凯伟
 * @date 2024/10/30
 */
public class ValueManager {
    private final Map<Value, TarValue> globalValueMap = new HashMap<>();
    private final Map<Value, TarValue> localValueMap = new HashMap<>();

    public TarValue getTargetValue(Value value) {
//        if (value instanceof ImmediateValue immediateValue) {
//            return new Immediate(immediateValue.getValue());
//        }
//        return globalValueMap.getOrDefault(value, localValueMap.getOrDefault(value, null));
        if (value instanceof ImmediateValue immediateValue) {
            return new Immediate(immediateValue.getValue());
        } else if (globalValueMap.containsKey(value)) {
            return globalValueMap.get(value);
        } else return localValueMap.getOrDefault(value, null);
    }

    public void putGlobal(Value value, Label label) {
        globalValueMap.put(value, label);
    }

    public int putLocals(FunctionIRType func) {
        return graphColoringManage(func);
    }

    private int basicManage(FunctionIRType func) {
        var registersName = List.of("s0", "s1", "s2", "s3", "s4", "s5", "s6", "s7");
        return manageMemory(func, registersName.stream().map(Register.REGISTER_MAP::get).toList(), this::refCountGlobalRegisterManage);
    }

    private int graphColoringManage(FunctionIRType func) {
        var registersName = List.of("s0", "s1", "s2", "s3", "s4", "s5", "s6", "s7");
        return manageMemory(func, registersName.stream().map(Register.REGISTER_MAP::get).toList(), this::graphColoringGlobalRegisterManage);
    }

    private void basicGlobalRegisterManage(List<Register> registers, List<AllocaInstruction> varInsts, FunctionIRType func) {
        Stack<Register> registersAllocator = new Stack<>();
        registersAllocator.addAll(registers);
//        for (var varInst : varInsts) {
//            if (registersAllocator.isEmpty()) {
//                break;
//            }
//            if (!varInst.getDataType().getArrayDims().isEmpty()) { // if is arrray
//                continue;
//            }
//
//            localValueMap.put(varInst, registersAllocator.pop());
//        }
        varInsts.stream()
                .filter(varInst -> !registersAllocator.isEmpty() && varInst.getDataType().getArrayDims().isEmpty())
                .forEach(varInst -> localValueMap.put(varInst, registersAllocator.pop()));
    }

    private void refCountGlobalRegisterManage(List<Register> registers, List<AllocaInstruction> varInsts, FunctionIRType func) {
        Stack<Register> registersAllocator = new Stack<>();
        registersAllocator.addAll(registers);

        var integerVarInsts = varInsts
                .stream()
                .filter(inst -> inst.getDataType().getArrayDims().isEmpty())
                .toList();

        Map<AllocaInstruction, Integer> refCounts = new HashMap<>();

//        for (var inst : integerVarInsts) {
//            int refCount = 0;
//            for (var use : inst.getUseList()) {
//                var loopWeight = ((Instruction) use.getUser()).getBasicBlock().getLoopNum();
//                if (loopWeight > 0) {
//                    refCount += 5 * loopWeight;
//                } else {
//                    refCount++;
//                }
//            }
//            refCounts.put(inst, refCount);
//        }
        integerVarInsts.forEach(inst -> {
            int refCount = 0;
            for (var use : inst.getUseList()) {
                var loopWeight = ((Instruction) use.getUser()).getBasicBlock().getLoopNum();
                refCount += loopWeight > 0 ? 5 * loopWeight : 1;
            }
            refCounts.put(inst, refCount);
        });
        var varInstsOrderByRefCount = new ArrayList<>(refCounts.entrySet().stream().sorted(Map.Entry.comparingByValue()).map(Map.Entry::getKey).toList());
        Collections.reverse(varInstsOrderByRefCount);

//        for (var inst : varInstsOrderByRefCount) {
//            if (registersAllocator.isEmpty()) {
//                break;
//            }
//            localValueMap.put(inst, registersAllocator.pop());
//        }
        varInstsOrderByRefCount.forEach(inst -> {
            if (registersAllocator.isEmpty()) {
                return;
            }
            localValueMap.put(inst, registersAllocator.pop());
        });
    }

    private void graphColoringGlobalRegisterManage(List<Register> registers, List<AllocaInstruction> varInsts, FunctionIRType func) {
        var liveVarAnalyzer = new LiveVariableAnalyzer(func);
        liveVarAnalyzer.analyze();

        varInsts = varInsts
                .stream()
                .filter(inst -> inst.getDataType().getArrayDims().isEmpty())
                .toList();

        var inSets = liveVarAnalyzer.getInSets();
        var outSets = liveVarAnalyzer.getOutSets();
        Map<BasicBlock, Set<AllocaInstruction>> activeSets = new HashMap<>();
//        for (var block : func.getBasicBlocks()) {
//            var activeSet = new HashSet<AllocaInstruction>();
//            activeSets.put(block, activeSet);
//            activeSet.addAll(inSets.get(block));
//            activeSet.addAll(outSets.get(block));
//        }
        func.getBasicBlocks().forEach(block -> {
            var activeSet = new HashSet<AllocaInstruction>();
            activeSets.put(block, activeSet);
            activeSet.addAll(inSets.get(block));
            activeSet.addAll(outSets.get(block));
        });
        var conflictDiagram = new ConflictDiagramBuilder(varInsts, liveVarAnalyzer.getDefSets(), activeSets).getConflictDiagram();

        var diagramForColor = conflictDiagram.copy();

        int degreeThreshold = registers.size();

        Stack<AllocaInstruction> nodesToColor = new Stack<>();

        while (!diagramForColor.isEmpty()) {
            AllocaInstruction candidate = null;
            for (var node : diagramForColor.getNodes()) {
                if (diagramForColor.getConflict(node).size() >= degreeThreshold) {
                    continue;
                }
                candidate = node;
                break;
            }

            if (candidate == null) {
                int maxDegree = degreeThreshold - 1;
                for (var node : diagramForColor.getNodes()) {
                    int degree = diagramForColor.getConflict(node).size();
                    if (diagramForColor.getConflict(node).size() <= maxDegree) {
                        continue;
                    }
                    maxDegree = degree;
                    candidate = node;
                    break;
                }
            } else {
                nodesToColor.push(candidate);
            }
            diagramForColor.removeNode(candidate);
        }

        while (!nodesToColor.isEmpty()) {
            var node = nodesToColor.pop();
            Set<Register> preservedRegs = new HashSet<>();
//            for (var conflictNode : conflictDiagram.getConflict(node)) {
//                if (localValueMap.containsKey(conflictNode) && localValueMap.get(conflictNode) instanceof Register registerInUse) {
//                    preservedRegs.add(registerInUse);
//                }
//            }
            conflictDiagram.getConflict(node).stream()
                    .filter(conflictNode -> localValueMap.containsKey(conflictNode) && localValueMap.get(conflictNode) instanceof Register)
                    .map(conflictNode -> (Register) localValueMap.get(conflictNode))
                    .forEach(preservedRegs::add);
//            for (var reg : registers) {
//                if (preservedRegs.contains(reg)) {
//                    continue;
//                }
//                localValueMap.put(node, reg);
//                break;
//            }
            registers.stream()
                    .filter(reg -> !preservedRegs.contains(reg))
                    .findFirst()
                    .ifPresent(reg -> localValueMap.put(node, reg));
        }
    }

    public List<Register> getRegistersInUse() {
        return localValueMap.values().stream().filter(elm -> elm instanceof Register).map(elm -> (Register) elm).distinct().toList();
    }

    public void clearLocals() {
        localValueMap.clear();
    }


    private int manageMemory(FunctionIRType func, List<Register> registersToAlloc, GlobalRegisterManager globalRegisterManager) {
        int argNum = func.getArguments().size();
        List<AllocaInstruction> allAllocaInsts = new ArrayList<>(func.getFirstBasicBlock().getInstructions()
                .stream()
                .filter(inst -> inst instanceof AllocaInstruction)
                .map(inst -> (AllocaInstruction) inst)
                .toList());

        List<AllocaInstruction> argAllocaInsts = allAllocaInsts.subList(0, argNum); // argn to arg0
        Collections.reverse(argAllocaInsts); // arg0 to argn
        List<AllocaInstruction> varAllocaInsts = allAllocaInsts.subList(argNum, allAllocaInsts.size());

//        for (int i = 0; i < argAllocaInsts.size() && i < 4; i++) {
//            var inst = argAllocaInsts.get(i);
//            localValueMap.put(inst, Register.REGISTER_MAP.get("a" + i));
//        }

        IntStream.range(0, Math.min(argAllocaInsts.size(), 4))
                .forEach(i -> localValueMap.put(argAllocaInsts.get(i), Register.REGISTER_MAP.get("a" + i)));
        globalRegisterManager.manageGlobalRegister(registersToAlloc, varAllocaInsts, func);

        int numOfArgOnRegister = Math.min(argNum, 4);
        int memoryRequire = getMemoryRequire(func, numOfArgOnRegister);

        int baseOffset = memoryRequire;
        var sp = Register.REGISTER_MAP.get("sp");

        for (int i = argAllocaInsts.size() - 1; i >= 4; i--) {
            var argInst = argAllocaInsts.get(i);
            baseOffset -= 4;
            localValueMap.put(argInst, new Offset(sp, baseOffset));
        }

        baseOffset -= numOfArgOnRegister * 4;

        for (var basicBlock : func.getBasicBlocks()) {
            for (var instruction : basicBlock.getInstructions()) {
                if ((instruction instanceof StoreInstruction)
                        || (instruction instanceof BrInstruction)
                        || (instruction instanceof ReturnInstruction)
                        || (instruction instanceof CallInstruction callInst && callInst.getType().getType() == IRType.VOID)) {
                    continue;
                }
                if (localValueMap.containsKey(instruction)) {
                    continue;
                }

                if (instruction instanceof AllocaInstruction allocaInst
                        && allocaInst.getDataType() instanceof ArrayIRType arrayIRType
                        && arrayIRType.getPtrNum() == 0) {
                    baseOffset -= 4 * arrayIRType.getTotalSize();
                } else {
                    baseOffset -= 4;
                }

                localValueMap.put(instruction, new Offset(sp, baseOffset));
            }
        }

        if (baseOffset != 0) {
            throw new RuntimeException();
        }

        return memoryRequire;
    }

    private int getMemoryRequire(FunctionIRType func, int numOfArgOnRegister) {
//        int memoryRequire = numOfArgOnRegister * 4;
//
//        for (var basicBlock : func.getBasicBlocks()) {
//            for (var instruction : basicBlock.getInstructions()) {
//                if ((instruction instanceof StoreInstruction)
//                        || (instruction instanceof BrInstruction)
//                        || (instruction instanceof ReturnInstruction)
//                        || (instruction instanceof CallInstruction callInst && callInst.getType().getType() == IRType.VOID)) {
//                    continue;
//                }
//                if (localValueMap.containsKey(instruction)) {
//                    continue;
//                }
//
//                if (instruction instanceof AllocaInstruction allocaInst
//                        && allocaInst.getDataType() instanceof ArrayIRType arrayIRType
//                        && arrayIRType.getPtrNum() == 0) {
//                    memoryRequire += 4 * arrayIRType.getTotalSize();
//                } else {
//                    memoryRequire += 4;
//                }
//            }
//        }
//        return memoryRequire;
        int memoryRequire = numOfArgOnRegister * 4;

        memoryRequire += func.getBasicBlocks().stream()
                .flatMap(basicBlock -> basicBlock.getInstructions().stream())
                .filter(instruction -> !(instruction instanceof StoreInstruction
                        || instruction instanceof BrInstruction
                        || instruction instanceof ReturnInstruction
                        || (instruction instanceof CallInstruction callInst && callInst.getType().getType() == IRType.VOID)))
                .filter(inst -> !localValueMap.containsKey(inst))
                .mapToInt(inst -> {
                    if (inst instanceof AllocaInstruction allocaInst
                            && allocaInst.getDataType() instanceof ArrayIRType arrayIRType
                            && arrayIRType.getPtrNum() == 0) {
                        return 4 * arrayIRType.getTotalSize();
                    } else {
                        return 4;
                    }
                })
                .sum();

        return memoryRequire;
    }
}

interface GlobalRegisterManager {
    void manageGlobalRegister(List<Register> registers, List<AllocaInstruction> varInsts, FunctionIRType func);
}